from stock.candlestick_charts import is_doji,is_dragonfly_doji,is_hammer,is_inverted_hammer,is_shooting_star,is_star,is_long_black_candle,is_medium_black_candle,is_long_white_candle,is_windy_wavey_candle
from stock.candlestick_charts import is_little_black_candle,is_little_white_candle,is_medium_white_candle
from stock.stock_base_daily import daily_directory,get_daily_path_name
import os
import pandas as pd
from utils import read_config,get_directory

# 读取配置信息
config = read_config()

# 涨幅
_increase = 7
# 上涨日
_rising_day = 'rising_day'
# 连续上涨日
_consecutive_rising_day = 'consecutive_rising_day'

file_name = "stock_analysis.csv"
_base_directory = config.get('settings','base_file')
_analysis_directory = config.get('settings','analysis_file')
base_path = get_directory(_base_directory,_analysis_directory)


def preprocessing_data(df,increase_percentage = _increase):
    # 初始化一个新列来追踪连续涨停的次数
    df[_consecutive_rising_day] = 0
    # 遍历DataFrame，更新大幅上涨次数
    for i in range(1, len(df)):
        _open = df.iloc[i]['open']
        _close = df.iloc[i]['close']
        if (_close - _open)/_open * 100 > increase_percentage :
            df.at[i, _consecutive_rising_day] = df.at[i-1, _consecutive_rising_day] + 1
        else:
            df.at[i, _consecutive_rising_day] = 0
    # 中阴线
    df['is_medium_black_candle'] = df.apply(is_medium_black_candle, axis=1)
    # 十字星
    df['is_doji'] = df.apply(is_doji, axis=1)
    # 长下影线（Dragonfly Doji）
    df['is_dragonfly_doji'] = df.apply(is_dragonfly_doji, axis=1)
    # 锤子线
    df['is_hammer'] = df.apply(is_hammer, axis=1)
    # 倒锤子线
    df['is_inverted_hammer'] = df.apply(is_inverted_hammer, axis=1)
    # 长上影线
    df['is_shooting_star'] = df.apply(is_shooting_star, axis=1)
    # 星线
    df['is_star'] = df.apply(is_star, axis=1)
    # 大阴线
    df['is_long_black_candle'] = df.apply(is_long_black_candle, axis=1)
    # 大阳线
    df['is_long_white_candle'] = df.apply(is_long_white_candle, axis=1)
    # 风高浪急线
    df['is_windy_wavey_candle'] = df.apply(is_windy_wavey_candle,axis=1)
    df['is_little_black_candle'] = df.apply(is_little_black_candle, axis=1)
    df['is_little_white_candle'] = df.apply(is_little_white_candle, axis=1)
    df['is_medium_white_candle'] = df.apply(is_medium_white_candle, axis=1)
    # 计算前一天的K线模式
    df['prev_medium_black_candle'] = df['is_medium_black_candle'].shift(1)
    df['prev_doji'] = df['is_doji'].shift(1)
    df['prev_dragonfly_doji'] = df['is_dragonfly_doji'].shift(1)
    df['prev_hammer'] = df['is_hammer'].shift(1)
    df['prev_inverted_hammer'] = df['is_inverted_hammer'].shift(1)
    df['prev_shooting_star'] = df['is_shooting_star'].shift(1)
    df['prev_star'] = df['is_star'].shift(1)
    df['prev_long_black_candle'] = df['is_long_black_candle'].shift(1)
    df['prev_long_white_candle'] = df['is_long_white_candle'].shift(1)
    df['prev_windy_wavey_candle'] = df['is_windy_wavey_candle'].shift(1)
    df['prev_little_black_candle'] = df['is_little_black_candle'].shift(1)
    df['prev_little_white_candle'] = df['is_little_white_candle'].shift(1)
    df['prev_medium_white_candle'] = df['is_medium_white_candle'].shift(1)
    df['pre_volume'] = df['volume'].shift(1)
    # 设置min_periods=5，确保即使窗口内的数据少于5行，也会进行计算
    df['trend'] = df['pctChg'].rolling(window=5).sum()
    # 设置min_periods=5，确保即使窗口内的数据少于5行，也会进行计算
    df['prev_trend'] = df['trend'].shift(1)
    # 计算后两天的最低价（不包含当前日）
    df['low_next_2_days'] = df['low'].rolling(window=2).min().shift(-2)
    # 计算后4天的最高价（不包含当前日）
    df['high_next_2_days'] = df['high'].rolling(window=2).max().shift(-2)
    # 计算后4天的最低价（不包含当前日）
    df['low_next_5_days'] = df['low'].rolling(window=5).min().shift(-5)
    # 计算后两天的最高价（不包含当前日）
    df['high_next_5_days'] = df['high'].rolling(window=5).max().shift(-5)
    # 计算后5天的最高收盘价（不包含当前日）
    df['high_close_next_5_days'] = df['close'].rolling(window=5).max().shift(-5)
    # 计算后5天的最低收盘价（不包含当前日）
    df['low_close_next_5_days'] = df['close'].rolling(window=5).min().shift(-5)
    # 计算后10天的最高收盘价（不包含当前日）
    df['high_close_next_10_days'] = df['close'].rolling(window=10).max().shift(-10)
    # 计算后10天的最低收盘价（不包含当前日）
    df['low_low_next_10_days'] = df['low'].rolling(window=10).min().shift(-10)
    return df


def preprocessing_all(increase_percentage=_increase,stock_codes=None):
    # 获取目录中的所有文件
    files = os.listdir(daily_directory)
    # 如果目录为空，直接返回
    if not files:
        print(f"目录 {daily_directory} 下没有文件.")
    # 遍历目录中的所有文件
    # 创建一个空的DataFrame
    rs = pd.DataFrame()
    if stock_codes is None:
        for filename in files:
            # 拼接文件的完整路径
            file_path = os.path.join(daily_directory, filename)
            ndf = _load_data_by_file(file_path, increase_percentage)
            rs = pd.concat([rs, ndf], ignore_index=True)
        # 假设df是您的DataFrame，已经包含了上述提到的列
    else:
        # 使用split方法分割字符串
        codes = stock_codes.split(',')
        # 遍历分割后的字符串列表
        for code in codes:
            # 移除可能存在的空白字符（如空格）
            code = code.strip()
            file_path = get_daily_path_name(code)
            ndf = _load_data_by_file(file_path, increase_percentage)
            rs = pd.concat([rs, ndf], ignore_index=True)
    return rs


def _load_data_by_file(file_path, increase_percentage):
    # 从CSV文件加载数据
    df = pd.read_csv(file_path, parse_dates=['date'])
    # 按日期排序
    df.sort_values('date', inplace=True)
    return preprocessing_data(df, increase_percentage)


# 把预处理的数据保存到file_path中
def preprocessing_by_all_file(file_name=file_name,stock_codes = None, increase_percentage = _increase):
    rs = preprocessing_all(increase_percentage = increase_percentage,stock_codes=stock_codes)
    __file = os.path.join(base_path, file_name)
    rs.to_csv(__file,index=False, encoding='utf-8-sig')


def load_file(file_name=file_name):
    __file = os.path.join(base_path, file_name)
    return pd.read_csv(__file, parse_dates=['date'])


if __name__=='__main__':
    preprocessing_by_all_file()
    print("完成！");